//
//  NSObjectHOMAdditions.m
//  HigherOrderMessaging
//
//  Created by Ofri Wolfus on 10/09/05.
//  Copyright 2005 Ofri Wolfus. All rights reserved.
//

#import "NSObjectHOMAdditions.h"
#import "HOMTrampoline.h"
#import "NSInvocationAdditions.h"
#import "HOMIteratedArgument.h"
#import "HOMUtilities.h"
#import "Message.h"
#import "HOMUtilities.h"
#include <objc/objc-class.h>
#include <objc/objc-runtime.h>
#include <malloc/malloc.h>


@implementation NSObject (NewHOM)

- (id)for:(int *)counter from:(int)start to:(int)end do:(Message *)msg {
	SEL				sel;
	marg_list		origArgs;
	marg_list		args;
	unsigned		argCount;
	Method			meth;
	unsigned		i;
	Class			iteratedArgCls;
	unsigned		count = 0U;
	char			*types;
	unsigned		argsSize;
	NSMutableArray	*result = nil;
	IMP				addObject = NULL;
	int				c;
	
	// Make sure we got a message
	DPAssert(msg != nil, @"The passed message must not be nil");
	
	// Init the variables we need
	sel = [msg selector];
	origArgs = [msg arguments];
	args = malloc(malloc_size([msg arguments]));
	argCount = hom_getNumberOfArguments(sel);
	iteratedArgCls = [HOMIteratedArgument class];
	types = calloc(argCount, sizeof(char));
	// If we don't implement this again as a class method, when -do: is being sent to a class, this will be called.
	// In that case we have to know if self is an instance or a class.
	meth = HOMObjectIsInstance(self) ? class_getInstanceMethod(self->isa, sel) : class_getClassMethod((Class)self, sel);
	
	// Make sure we respond to the passed message
	DPAssert(meth != NULL, @"+++ %s[%@ %@]: selector not recognized", HOMObjectIsInstance(self) ? "-" : "+", [self className], _cmd);
	
	// Get the size of our arguments. We don't use the size passed with our message object as it may be larger then the real size.
	argsSize = method_getSizeOfArguments(meth);
	
	// First, copy the arguments list since we're going to modify it
	memcpy(args, origArgs, malloc_size(args));
	
	// Scan the arguments of our message and find instances of HOMIteratedArgument.
	// Our types array will mark the places of these with the letter 'z' so we don't have to check again.
	// This loop also finds the largest number of arguments that needs to be iterated, and stores it in count.
	for (i = 2U; i < argCount; i++) {
		types[i] = *hom_getArgumentTypeAtIndex(meth, i);
		if (types[i] == *@encode(id)) {
			id	obj = marg_getValue(args, hom_getArgumentOffsetAtIndex(meth, i), id);
			
			if ([obj isKindOfClass:iteratedArgCls] && [obj count] > count) {
				types[i] = 'z';	// This is our special mark for an iterated argument
				count = [obj count];
			}
		}
	}
	
	// Allocate an array for the results if the return type is id or Class.
	if ((*(method_getTypeEncoding(meth)) == _C_ID || *(method_getTypeEncoding(meth)) == _C_CLASS)) {
		result = [[[NSMutableArray alloc] initWithCapacity:count * (end - start)] autorelease];
		addObject = [result methodForSelector:@selector(addObject:)];
	}
	
	// Send the passed message as many times as we've been told to
	for (c = start; c < end; c++) {
		// Update their counter
		if (counter)
			*counter = c;
		
		// Iterated arguments are used
		if (count > 0U) {
			// Copy the types array 'cause we're going to modify it inside the loop
			char *types_copy = calloc(argCount, sizeof(char));
			memcpy(types_copy, types, argCount * sizeof(char));
			
			// Loop through all objects in all our iterated arguments, update our args list, and resend it
			for (i = 0U; i < count; i++) {
				unsigned j;
				
				// Update the remaining arguments in our list
				for (j = 2U; j < argCount; j++) {
					if (types[j] == 'z') {
						int offset = hom_getArgumentOffsetAtIndex(meth, j);
						HOMIteratedArgument *arg = marg_getValue(origArgs, offset, HOMIteratedArgument *);
						
						// Reset the iterated arguments at the beginning of the loop
						if (i == 0)
							[arg reset];
						
						// Set the argument in our ars frame to the next argument in the iterated set
						marg_setValue(args, offset, id, [arg nextObject]);
						
						// When there are no arguments left for this argument iterator, we mark its type as id
						// so we'll skip it in the next run.
						if ([arg isFinished])
							types[j] = _C_ID;
					}
				}
				
				// Send the updated arguments list. If result is nil it means the result can't be stored in an array.
				result ? addObject(result, @selector(addObject:), objc_msgSendv(self, sel, argsSize, args))
					   : objc_msgSendv(self, sel, argsSize, args);
			}
			
			// Free the types array we copied at the beginning
			free(types_copy);
		} else {	// There are no iterated arguments so we simply send our message
			if (*(method_getTypeEncoding(meth)) == _C_ID || *(method_getTypeEncoding(meth)) == _C_CLASS)
				[result addObject:objc_msgSendv(self, sel, argsSize, origArgs)];
		}
	}

	// Clean up
	free(types);
	free(args);

	if (count == 0U && [result count] == 1U)
		result = [result objectAtIndex:0U];
	
	return result ?: self;
}

- (id)do:(Message *)msg {
	return [self for:NULL from:0 to:1 do:msg];
}

- (id)repeatOf:(Message *)msg for:(int)xTimes {
	return [self for:NULL from:0 to:xTimes do:msg];
}

- (id)receive:(Message *)msg {
	DPAssert(msg != nil, @"The passed message must not be nil");
	SEL sel = [msg selector];
	
	if ([self respondsToSelector:sel])
		return objc_msgSendv(self, sel, [msg argumentsSize], [msg arguments]);
	
	return nil;
}

@end
